
import cv2
import numpy as np
import torch
import torch.utils.data as Data
from sklearn import metrics

from model import RAS
from utils import trans_im, trans_gt

START_ID = 1
END_ID = 1
IMAGES_NUM = END_ID - START_ID + 1
BATCH_SIZE = 1


def get_val_data(start_image_id, end_image_id):
    image_num = end_image_id - start_image_id + 1  # 包含 end_id
    x = np.zeros([image_num, 3, 500, 500])
    y = np.zeros([image_num, 1, 500, 500])
    im_ids = np.zeros([image_num, 1])
    for i in range(image_num):
        img_id = start_image_id + i
        im_path = f"data/train/images/{img_id:04d}.jpg"
        gt_path = f"data/train/ground_truth_mask/{img_id:04d}.png"
        x[i, :, :, :] = trans_im(im_path)
        y[i, :, :, :] = trans_gt(gt_path)
        im_ids[i] = img_id
    x = torch.FloatTensor(x)
    y = torch.FloatTensor(y)
    im_ids = torch.Tensor(im_ids)
    torch_dataset = Data.TensorDataset(x, y, im_ids)
    loader = Data.DataLoader(
        dataset=torch_dataset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=4,
    )
    return loader


if __name__ == "__main__":
    ras = RAS()
    device = torch.device("cpu")
    ras.to(device)
    ras.load_state_dict(torch.load("data/model/epoch_199_params.pkl"))

    Y_test = []
    Y_prob = []
    loader = get_val_data(START_ID, END_ID)
    for step, (batch_x, batch_y, im_id) in enumerate(loader):
        img_id = int(im_id.numpy()[0, 0])
        im_path_pre = f"data/visualization/{img_id:04d}"
        batch_y_prob = ras.test(batch_x.to(device), im_path_pre)
        Y_test.append(batch_y.numpy().flatten().astype(np.int32))
        Y_prob.append(batch_y_prob.cpu().numpy().flatten())
        im_path = f"data/visualization/{img_id:04d}.png"
        im = Y_prob[-1].reshape(500, 500, 1) * 255
        cv2.imwrite(im_path, im.astype(np.uint8))
        if step % 20 == 0:
            print(f"finished step {step}")
    auc = metrics.roc_auc_score(
        np.array(Y_test).flatten(), np.array(Y_prob).flatten())
    print(f"auc is {auc}")

    MIOU = 0.
    for i in range(START_ID, END_ID + 1):
        y = cv2.imread(f"data/visualization/{i:04d}.png")
        gt = cv2.imread(f"data/train/ground_truth_mask/{i:04d}.png")
        y = cv2.resize(y, (gt.shape[1], gt.shape[0]), interpolation=cv2.INTER_AREA)
        y = y[:, :, 0]
        gt = gt[:, :, 0]
        y_t = np.sum(y > 255. / 2)
        g_t = np.sum(gt > 255. / 2)
        tp = np.sum(gt[y > 255. / 2] > 255. / 2)
        if (y_t + g_t - tp) != 0:
            MIOU += tp / (y_t + g_t - tp)
    print(f"MIOU is {MIOU / IMAGES_NUM}")